from __future__ import annotations

"""
AgentManager
------------
常驻后台循环，统一调度 Agent 策略：
- 定期从数据库读取策略配置（agent_strategies）
- 根据 enabled 状态创建/停止对应的 AgentWorker
- 每个 AgentWorker 按 interval_seconds 触发一次“心跳”逻辑（当前仅记录日志与心跳时间）

暂未接入大模型，先跑通调度与生命周期管理。
"""

import asyncio
import logging
from typing import Dict, Optional

from sqlalchemy import select, update

from ..db import SessionLocal
from ..models import AgentStrategy, utcnow
from .ai_autotrader import AIAutoTrader, AIAutoTraderConfig

logger = logging.getLogger(__name__)


class AgentWorker:
    """单个策略的运行实例"""

    def __init__(self, strategy: AgentStrategy):
        self.strategy = strategy
        self.updated_at = strategy.updated_at
        self._task: Optional[asyncio.Task] = None
        self._stop_event = asyncio.Event()
        self._wake_event = asyncio.Event()
        self._trader: Optional[AIAutoTrader] = None

    async def start(self) -> None:
        if self._task and not self._task.done():
            return
        self._stop_event.clear()
        await self._start_trader()
        self._task = asyncio.create_task(self._loop())
        logger.info("[AgentWorker] started id=%s name=%s", self.strategy.id, self.strategy.name)

    async def stop(self) -> None:
        self._stop_event.set()
        self._wake_event.set()
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        if self._trader:
            await self._trader.stop()
            self._trader = None
        logger.info("[AgentWorker] stopped id=%s name=%s", self.strategy.id, self.strategy.name)

    def refresh(self, strategy: AgentStrategy) -> None:
        """更新策略配置并唤醒"""
        self.strategy = strategy
        self.updated_at = strategy.updated_at
        self._wake_event.set()

    async def _loop(self) -> None:
        while not self._stop_event.is_set():
            try:
                await self._run_once()
            except Exception as exc:  # noqa: BLE001
                logger.error("[AgentWorker] run_once error id=%s: %s", self.strategy.id, exc, exc_info=True)

            interval = max(self.strategy.interval_seconds, 1)
            try:
                await asyncio.wait_for(self._wake_event.wait(), timeout=interval)
            except asyncio.TimeoutError:
                pass
            self._wake_event.clear()

    async def _run_once(self) -> None:
        # 仅更新心跳；自动交易在 AIAutoTrader 内部循环运行
        async with SessionLocal() as session:
            stmt = (
                update(AgentStrategy)
                .where(AgentStrategy.id == self.strategy.id)
                .values(last_heartbeat=utcnow())
            )
            await session.execute(stmt)
            await session.commit()

    async def _start_trader(self) -> None:
        """按照 AgentStrategy 配置启动 AIAutoTrader。"""
        if self.strategy.id is None:
            raise ValueError("AgentStrategy id is required to start trader")

        run_id = f"agent-{self.strategy.id}"
        cfg = AIAutoTraderConfig(
            run_id=run_id,
            account_name=self.strategy.account_name,
            connector_name=self.strategy.connector_name,
            trading_pair=self.strategy.trading_pair,
            max_notional_usd=self.strategy.max_notional_usd,
            max_slippage_bps=self.strategy.max_slippage_bps,
            max_daily_loss_usd=self.strategy.max_daily_loss_usd,
            mode=self.strategy.mode,
            interval_seconds=self.strategy.interval_seconds,
        )
        # 如果已有 trader，先停止
        if self._trader:
            await self._trader.stop()
        self._trader = AIAutoTrader(cfg)
        await self._trader.start()


class AgentManager:
    """统一管理 AgentWorker 的生命周期"""

    def __init__(self, poll_interval: int = 5):
        self.poll_interval = poll_interval
        self._task: Optional[asyncio.Task] = None
        self._stop_event = asyncio.Event()
        self._refresh_event = asyncio.Event()
        self._workers: Dict[int, AgentWorker] = {}

    async def start(self) -> None:
        if self._task and not self._task.done():
            return
        self._stop_event.clear()
        self._task = asyncio.create_task(self._loop())
        logger.info("[AgentManager] started")

    async def stop(self) -> None:
        self._stop_event.set()
        self._refresh_event.set()
        if self._task:
            await self._task
        for worker in list(self._workers.values()):
            await worker.stop()
        self._workers.clear()
        logger.info("[AgentManager] stopped")

    def request_refresh(self) -> None:
        """外部调用：策略变更后立即触发一次同步"""
        self._refresh_event.set()

    def runtime_status(self) -> Dict[int, str]:
        """返回内存中各策略的运行状态"""
        out: Dict[int, str] = {}
        for sid, worker in self._workers.items():
            running = worker._task is not None and not worker._task.done()  # noqa: SLF001
            out[sid] = "running" if running else "stopped"
        return out

    async def _loop(self) -> None:
        while not self._stop_event.is_set():
            await self._sync_strategies()
            try:
                await asyncio.wait_for(self._refresh_event.wait(), timeout=self.poll_interval)
            except asyncio.TimeoutError:
                pass
            self._refresh_event.clear()

    async def _sync_strategies(self) -> None:
        async with SessionLocal() as session:
            result = await session.execute(select(AgentStrategy))
            strategies = result.scalars().all()

        active_ids = set()
        for strategy in strategies:
            if not strategy.enabled:
                # 如果 DB 关闭了策略，停止对应 worker
                if strategy.id in self._workers:
                    await self._workers[strategy.id].stop()
                    self._workers.pop(strategy.id, None)
                continue

            if strategy.id is None:
                continue

            active_ids.add(strategy.id)
            worker = self._workers.get(strategy.id)
            if worker is None:
                worker = AgentWorker(strategy)
                self._workers[strategy.id] = worker
                await worker.start()
                continue

            # 如果策略配置有更新，则刷新并唤醒
            if worker.updated_at != strategy.updated_at:
                await worker.stop()
                worker = AgentWorker(strategy)
                self._workers[strategy.id] = worker
                await worker.start()

        # 清理已删除的策略
        for sid in list(self._workers.keys()):
            if sid not in active_ids:
                await self._workers[sid].stop()
                self._workers.pop(sid, None)


# 供外部直接引用的单例
agent_manager = AgentManager()
